/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 *
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include <vector>
#include <cstdint>
#include <cmath>
#include <random>
#include <iostream>
#include <algorithm>
#include "acl/acl.h"
#include "elewise/elewise_host.h"
#include "elewise/elewise_device.h"
#include "common/example_common.h"

namespace {
// 描述算子的输入输出以及临时计算资源
// 注意：第二个输入用 int，以匹配代码生成（你报错日志里看到的是 int 而不是 int32_t）
using PowOpTraits = ATVC::OpTraits<
    ATVC::OpInputs<float, int>,   // 输入: base(float), exp(int)
    ATVC::OpOutputs<float>,       // 输出: out(float)
    ATVC::OpTemps<float, float>   // 临时buffer
>;

// 传入编译态参数ATVC::OpTraits
template<class Traits>
struct PowComputeFunc {
    template<typename T, typename U>
    __aicore__ inline void operator()(AscendC::LocalTensor<T> base,
                                      AscendC::LocalTensor<int> expInt,
                                      AscendC::LocalTensor<T> out,
                                      AscendC::LocalTensor<U> tempBuffer1,
                                      AscendC::LocalTensor<U> tempBuffer2)
    {
        uint32_t tiledCnt = out.GetSize();

        // tempBuffer1 = ln(base)
        AscendC::Ln(tempBuffer1, base, tiledCnt);

        // tempBuffer2 = float(expInt)，需要指定 RoundMode
        AscendC::Cast(tempBuffer2, expInt, AscendC::RoundMode::CAST_NONE, tiledCnt);

        // tempBuffer1 = expInt * ln(base)
        AscendC::Mul(tempBuffer1, tempBuffer1, tempBuffer2, tiledCnt);

        // out = exp(tempBuffer1) = pow(base, expInt)
        AscendC::Exp(out, tempBuffer1, tiledCnt);
    }
};


// 生成输入和golden
void InitializeData(int32_t eleNum,
                    std::vector<float> &inputX,
                    std::vector<int> &inputY,
                    std::vector<float> &golden)
{
    std::random_device rd;
    std::mt19937 gen(rd());
    // base > 0，避免 ln 非法
    std::uniform_real_distribution<float> disBase(1.0f, 5.0f);
    std::uniform_int_distribution<int> disExp(0, 5);

    for (int i = 0; i < eleNum; ++i) {
        inputX[i] = disBase(gen);
        inputY[i] = disExp(gen);
        golden[i] = std::pow(inputX[i], static_cast<float>(inputY[i]));
    }
}
} // namespace

/*
 * 该函数为PowCustom算子核函数入口
 * base     Device上的gm地址，指向Pow算子底数
 * exp      Device上的gm地址，指向Pow算子指数
 * out      Device上的gm地址，指向Pow算子输出
 * param    Device上的gm地址，指向运行态ATVC::EleWiseParam数据
*/
template<class OpTraits>
__global__ __aicore__ void PowCustom(GM_ADDR base, GM_ADDR exp, GM_ADDR out, ATVC::EleWiseParam param)
{
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    auto op = ATVC::Kernel::EleWiseOpTemplate<PowComputeFunc<OpTraits>>();
    // ✅ 参数顺序必须是：所有输入 -> 所有输出 -> &param
    op.Run(base, exp, out, &param);
}

int main()
{
    // init data
    int32_t eleNum = 8 * 2048;
    size_t inputXByteSize = static_cast<size_t>(eleNum) * sizeof(float);   // base: float
    size_t inputYByteSize = static_cast<size_t>(eleNum) * sizeof(int);     // exp:  int
    size_t outputByteSize = static_cast<size_t>(eleNum) * sizeof(float);   // out:  float

    std::vector<float> inputX(eleNum);
    std::vector<int>   inputY(eleNum);
    std::vector<float> golden(eleNum);
    InitializeData(eleNum, inputX, inputY, golden);

    ATVC::EleWiseParam param;

    // 计算Pow算子的运行态参数param
    if (!ATVC::Host::CalcEleWiseTiling<PowOpTraits>(eleNum, param)) {
        printf("Elewise tiling error.");
        return -1;
    }

    // 初始化Acl资源与数据
    aclrtContext context;
    aclrtStream stream = nullptr;
    int32_t deviceId = 0;
    InitializeACL(context, stream, deviceId);

    uint8_t *xDevice;
    uint8_t *yDevice;
    uint8_t *outDevice;
    uint8_t *outHost;

    // 分配并拷贝输入X (float)
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputXByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemcpy(xDevice, inputXByteSize, inputX.data(), inputXByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // 分配并拷贝输入Y (int)
    CHECK_ACL(aclrtMalloc((void **)&yDevice, inputYByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemcpy(yDevice, inputYByteSize, inputY.data(), inputYByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // 分配输出
    CHECK_ACL(aclrtMallocHost((void **)(&outHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&outDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    // 调用自定义的Kernel API
    PowCustom<PowOpTraits><<<param.tilingData.blockNum, nullptr, stream>>>(xDevice, yDevice, outDevice, param);

    // 同步，拷回结果
    CHECK_ACL(aclrtSynchronizeStream(stream));
    CHECK_ACL(aclrtMemcpy(outHost, outputByteSize, outDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    std::vector<float> outputY(reinterpret_cast<float*>(outHost),
                               reinterpret_cast<float*>(outHost) + eleNum);

    // 释放资源
    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(outDevice));
    CHECK_ACL(aclrtFreeHost(outHost));

    CleanACL(stream, context, deviceId);

    // 精度验证
    if (!VerifyResults(golden, outputY)) {
        return -1;
    }
    printf("Accuracy verification passed.\n");
    return 0;
}
